import torch
import math


def init_weights(m):
#    print(m)

    if type(m) == torch.nn.Linear:
        m.weight.data.fill_(1.0/48)
        m.bias.data.fill_(0.01)
#        print(m.weight)
#        print(m.bias)
    elif type(m) == torch.nn.Conv2d:
        torch.nn.init.xavier_uniform_(m.weight)
#        print(m.weight)

# Creating a CNN class
class ConvNeuralNet(torch.nn.Module):
    #  Determine what layers and their order in CNN object
    def __init__(self, numHops, numClasses):
        super(ConvNeuralNet, self).__init__()

        timePoolSize = math.ceil(numHops / 8)

        dropoutProb = 0.2
        NumF = 12

        self.model = torch.nn.Sequential(
            # imageInputLayer([numHops,afe.FeatureVectorLength]),

            torch.nn.LayerNorm([98,50]),
            # NumF,98x50
            torch.nn.Conv2d(in_channels=1, out_channels=NumF, kernel_size=3, padding_mode='replicate', stride=1, padding=1),
            # 98x50
            torch.nn.BatchNorm2d(num_features=NumF),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1 ,stride=2),
            # 2*NumF,49x25
            torch.nn.Conv2d(in_channels=NumF, out_channels=2 * NumF, kernel_size=3, padding_mode='replicate', stride=1,
                            padding=1),
            # 49x25
            torch.nn.BatchNorm2d(num_features=2 * NumF),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1,stride=2),

            # 4*NumF,25x13
            torch.nn.Conv2d(in_channels=2 * NumF, out_channels=4 * NumF, kernel_size=3, padding_mode='replicate', stride=(1,2),
                            padding=1),
            # 25x7
            torch.nn.BatchNorm2d(num_features=4 * NumF),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=3, padding=1,stride=2),
            # 4*NumF,13x4
            torch.nn.Conv2d(in_channels=4 * NumF, out_channels=4 * NumF, kernel_size=3, padding_mode='replicate', stride=(1,2),
                            padding=1),
            # 13x2
            torch.nn.BatchNorm2d(num_features=4 * NumF),
            torch.nn.ReLU(),

            # 4*NumF,13x2,
            torch.nn.Conv2d(in_channels=4 * NumF, out_channels=4 * NumF, kernel_size=3, padding_mode='replicate', stride=(1,2),
                            padding=1),
            # 13x1
            torch.nn.BatchNorm2d(num_features=4 * NumF),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=[timePoolSize, 1]),
            torch.nn.Dropout(dropoutProb),
            torch.nn.Flatten(start_dim=1),
            # numClasses
            # size of each input sample,size of each output sample
            torch.nn.Linear(in_features=4 * NumF, out_features=numClasses),
            
            # classificationLayer(Classes=classes,ClassWeights=classWeights)];

        )

    def init(self):
        self.model.apply(init_weights)

    # Progresses data across layers
    def forward(self, x):
        return self.model(x)


